#include "analyzer.h"

#include <algorithm>
#include <functional>
#include <iostream>
#include <ostream>
#include <type_traits>
#include <utility>
#include <vector>

#include "arithmetic.h"
#include "ast.h"
#include "semantics.h"
#include "symbol.h"
#include "token.h"
#include "types.h"
#include "utils.h"

#define DECL_VISIT(className) \
    ASTNode* SemanticAnalyzer::Visit##className(className* node)

using namespace Utils;
using namespace AST;
using namespace Semantics;
using Context = SymbolMap::Context;

void SemanticAnalyzer::Cast(ASTNode*& origin, Type::Basic to) {
    auto cast = new CastExpr();
    cast->from = origin->As<Expression>()->info->baseType;
    cast->to = to;
    cast->origin = origin;
    origin = cast;
    Visit(origin);
}

Type::Basic SemanticAnalyzer::HandleImplicitCast(ASTNode*& lhs, ASTNode*& rhs) {
    Type::Basic tl = lhs->As<Expression>()->info->baseType;
    Type::Basic tr = rhs->As<Expression>()->info->baseType;
    if (tl == tr) return tl;
    bool swap = tl == Type::Basic::Int;
    // cast rhs from int to float
    if (swap) std::swap(lhs, rhs);
    Cast(rhs, Type::Basic::Float);
    if (swap) std::swap(lhs, rhs);
    return Type::Basic::Float;
}

DECL_VISIT(AssignExpr) {
    Visit(node->lhs);
    Visit(node->rhs);

    // array = initializer
    if (node->lhs->As<IdentifierExpr>() && node->rhs->As<InitializerExpr>()) {
        auto name = node->lhs->As<IdentifierExpr>()->name;
        auto* init = node->rhs->As<InitializerExpr>();
        auto* block = new Block();
        for (auto [pos, val] : init->parsed) {
            auto* e = new AddrExpr();
            e->base = new IdentifierExpr(name);
            e->offset = new LiteralExpr(pos);
            auto* ins = new AssignExpr();
            ins->lhs = e;
            ins->rhs = val;
            block->Add(ins);
        }
        Visit((ASTNode*&)block);
        return block;
    }

    node->info = node->lhs->As<Expression>()->info;
    Assert(node->info->As<LValue>(),
           "the left operand of '=' should be a left value");
    auto rhs = node->rhs->As<Expression>();
    if (rhs->info->baseType != node->info->baseType) {
        auto cast = new CastExpr();
        cast->origin = rhs;
        cast->from = rhs->info->baseType;
        cast->to = node->info->baseType;
        node->rhs = cast;
        Visit(node->rhs);
    }
    return node;
}

DECL_VISIT(CallExpr) {
    Visit(node->func);
    Visit(node->args);

    // parentheses
    if (node->func->As<VoidExpr>()) {
        return node->args;
    }

    // function
    auto func = node->func->As<IdentifierExpr>();
    Assert(func, "invalid '('");

    auto funcInfo = func->info->As<Function>();
    Assert(funcInfo, "expect a function before '('");

    // generate argument list
    auto args = node->args->As<Expression>();
    auto argList = new ArgList();
    if (!args->As<VoidExpr>()) {
        while (true) {
            auto cur = args->As<OperatorExpr>();
            if (!cur || cur->op != ArithOp::Comma) break;
            argList->Add(cur->rhs);
            args = cur->lhs->As<Expression>();
        }
        argList->Add(args);
    }
    reverse(argList->items.begin(), argList->items.end());

    // argument type casting
    if (!funcInfo->hasVA) {
        Assert(argList->items.size() == funcInfo->argType.size(),
               "argument size error");
        for (size_t i = 0; i < argList->items.size(); ++i) {
            auto& arg = argList->items[i];
            auto argInfo = arg->As<Expression>()->info;
            auto dest = funcInfo->argType[i];
            if (!dest.isPtr) {
                Assert(argInfo->As<Value>(), "argument type error");
                Cast(arg, dest.base);
            } else {
                auto argPtr = argInfo->As<Pointer>();
                Assert(argPtr && argPtr->baseType == dest.base,
                       "argument type error");
            }
        }
    }

    node->args = argList;
    node->info = new Value(funcInfo->baseType);
    return node;
}

DECL_VISIT(ArrayExpr) {
    Visit(node->lhs);
    Visit(node->rhs);

    Assert(node->rhs->As<Expression>()->info->baseType == Type::Basic::Int,
           "array index should be an integer");

    auto* lval = node->lhs->As<Expression>()->info->As<Pointer>();
    if (!lval) Err("expect pointer type");
    auto shape = lval->shape;
    int len = 1;
    for (auto x : shape) len *= x;

    if (shape.empty()) {
        // one dimension
        node->info = new LValue(lval->baseType);
    } else {
        // multi-dimension
        shape.pop_back();
        node->info = new Pointer(lval->baseType, shape);
    }

    // simplify addressing
    auto* delta = new OperatorExpr(ArithOp::Mul);
    delta->lhs = node->rhs;
    delta->rhs = new LiteralExpr(len);

    AddrExpr* replaced;
    if (node->lhs->As<IdentifierExpr>()) {
        replaced = new AddrExpr();
        replaced->base = node->lhs;
        replaced->offset = delta;
    } else if (node->lhs->As<AddrExpr>()) {
        replaced = (AddrExpr*)node->lhs;
        auto* offset = new OperatorExpr(ArithOp::Add);
        offset->lhs = replaced->offset;
        offset->rhs = delta;
        replaced->offset = offset;
    } else {
        Err("");
    }
    replaced->info = node->info;
    Visit(replaced->offset);

    // get init value of global variable
    if (symbolMap.GetContext() == Context::Global) {
        auto baseInfo = replaced->base->As<Expression>()->info->As<Pointer>();
        auto offInfo = replaced->offset->As<Expression>()->info->As<Constant>();
        Assert(baseInfo, "expect an array identifier");
        Assert(offInfo, "expect an constant as array index");
        ArithValue init(0);
        if (baseInfo->init) {
            auto it = baseInfo->init->find(offInfo->value.v.i);
            if (it != baseInfo->init->end()) init = ArithValue(it->second);
        }
        ASTNode* replaced = new LiteralExpr(init);
        Visit(replaced);
        return replaced;
    }

    return replaced;
}

DECL_VISIT(OperatorExpr) {
    Visit(node->lhs);
    Visit(node->rhs);

    auto* lhs = node->lhs->As<Expression>();
    auto* rhs = node->rhs->As<Expression>();
    switch (node->op) {
        case ArithOp::Comma:
            node->info = rhs->info;
            break;
        case ArithOp::Add:
        case ArithOp::Sub:
        case ArithOp::Not:
            if (!rhs->info->As<Value>()) {
                Err("expect value after operator");
            }
            if (lhs->info->As<Void>()) {
                auto type = ArithOp::GetType(rhs->info->baseType, node->op);
                node->info = new Value(type);
                break;
            }
        default:
            if (!lhs->info->As<Value>() || !rhs->info->As<Value>()) {
                Err("expect values around operator");
            }
            auto type = HandleImplicitCast(node->lhs, node->rhs);
            type = ArithOp::GetType(type, node->op);
            node->info = new Value(type);
            break;
    }

    // constant folding
    if (LiteralExpr* lhs = node->lhs->As<LiteralExpr>(),
        *rhs = node->rhs->As<LiteralExpr>();
        lhs && rhs && node->op != ArithOp::Comma) {
        ArithValue ret = ArithOp::Calc(lhs->value, rhs->value, node->op);
        ASTNode* replaced = new LiteralExpr(ret);
        Visit(replaced);
        return replaced;
    }
    if (auto rhs = node->rhs->As<LiteralExpr>();
        node->lhs->As<VoidExpr>() && rhs) {
        ArithValue ret = ArithOp::Calc(rhs->value, node->op);
        ASTNode* replaced = new LiteralExpr(ret);
        Visit(replaced);
        return replaced;
    }
    return node;
}

DECL_VISIT(TernaryExpr) {
    Visit(node->cond);
    Visit(node->tVal);
    Visit(node->fVal);

    auto type = HandleImplicitCast(node->tVal, node->fVal);
    node->info = new Value(type);
    return node;
}

DECL_VISIT(LiteralExpr) {
    node->info = new Constant(node->value);
    return node;
}

DECL_VISIT(VoidExpr) {
    node->info = new Void();
    return node;
}

DECL_VISIT(IdentifierExpr) {
    auto symbol = symbolMap.Find(node->name);

    if (auto info = symbol->info->As<Constant>()) {
        ASTNode* replaced = new LiteralExpr(info->value);
        Visit(replaced);
        return replaced;
    }

    // get init value of global variable
    if (symbolMap.GetContext() == Context::Global) {
        if (auto info = symbol->info->As<LValue>()) {
            ASTNode* replaced = new LiteralExpr(info->init);
            Visit(replaced);
            return replaced;
        }
    }

    node->info = symbol->info;
    return node;
}

DECL_VISIT(InitializerExpr) { return node; }

DECL_VISIT(CastExpr) {
    Visit(node->origin);
    auto origin = node->origin->As<Expression>();
    Assert(origin->info->baseType == node->from, "original type error");
    Assert(node->from != Type::Basic::Void, "unexpected void type");
    if (node->from == node->to) return node->origin;

    if (auto e = origin->As<LiteralExpr>()) {
        ASTNode* ret;
        if (node->to == Type::Basic::Int) {
            Assert(node->from == Type::Basic::Float, "unexpected type");
            ret = new LiteralExpr((int)e->value.v.f);
        } else if (node->to == Type::Basic::Float) {
            Assert(node->from == Type::Basic::Int, "unexpected type");
            ret = new LiteralExpr((float)e->value.v.i);
        } else
            Err("unexpected type");
        Visit(ret);
        return ret;
    }

    node->info = new Value(node->to);
    return node;
}

DECL_VISIT(AddrExpr) {
    Visit(node->base);
    Visit(node->offset);
    auto base = node->base->As<Expression>();
    node->info = new LValue(base->info->baseType);
    return node;
}

DECL_VISIT(ArgList) {
    Err("");
    return node;
}

DECL_VISIT(Block) {
    block.emplace_back(node);
    if (symbolMap.GetContext() == Context::None) {
        symbolMap.CreateEnv(Context::Global);
    } else {
        symbolMap.CreateEnv(Context::Local);
    }
    std::vector<ASTNode*> items;
    for (auto& item : node->items) {
        Visit(item);
        if (!item->As<VoidExpr>()) items.push_back(item);
    }
    node->items = items;
    symbolMap.ExitEnv();
    block.pop_back();
    return node;
}

// flatten initializer list:
// int a[2][2] = {{1, 2}, {3, 4}}
//            => {1, 2, 3, 4}
void FlattenInitList(VarDecl* node, const std::vector<int>& shape) {
    auto init = node->init->As<InitializerExpr>();
    if (!init) return;

    std::vector<std::pair<int, ASTNode*>> flatten;
    std::function<int(ASTNode*)> scan = [&](ASTNode* cur) {
        auto init = cur->As<InitializerExpr>();
        if (!init) {
            flatten.emplace_back(0, cur);
            return 0;
        }
        int maxDep = -1;
        for (auto& item : init->items) {
            maxDep = std::max(maxDep, scan(item) + 1);
        }
        if (maxDep == -1)
            flatten.emplace_back(1, nullptr);
        else
            flatten.back().first = maxDep;
        return flatten.back().first;
    };

    int depth = scan(node->init);
    Assert(depth <= shape.size(), "too many braces in initializer list");

    std::vector<int> diff(shape.size() + 1);
    diff[0] = 1;
    for (int i = 0; i < shape.size(); ++i) diff[i + 1] = diff[i] * shape[i];
    int pos = 0;

    for (auto [dep, item] : flatten) {
        Assert(pos <= diff.back(), "initializer list too large");
        if (item) init->parsed.emplace_back(pos, item);
        int d = diff[dep];
        pos = pos / d * d + d;
    }
    init->items.clear();
}

DECL_VISIT(VarDecl) {
    std::vector<int> shape;
    // get array shape
    for (auto& x : node->shape) {
        Visit(x);
        if (auto e = x->As<LiteralExpr>()) {
            Assert(e->value.type == Type::Basic::Int,
                   "array index type should be int");
            shape.push_back(e->value.v.i);
        } else if (x->As<VoidExpr>()) {
            shape.push_back(0);
        } else
            Err("expect constant in array index");
    }

    if (node->init) {
        // handle initializer list
        FlattenInitList(node, shape);
        Visit(node->init);
    }

    // constant
    if (shape.empty() && node->isConst) {
        Cast(node->init, node->baseType);
        auto init = node->init->As<LiteralExpr>();
        Assert(init, "expect constant in const declaration");
        auto symbol = Symbol::CreateConstant(node->name, init->value);
        symbolMap.Add(symbol);
        // ignore declaration of constant
        return new VoidExpr();
    }

    // pointer
    if (!shape.empty() && symbolMap.GetContext() == Context::Argument) {
        shape.pop_back();
        auto symbol = Symbol::CreateArray(node->baseType, node->name, shape, 0);
        symbolMap.Add(symbol);

        auto replaced = new IdentifierExpr(node->name);
        Visit((ASTNode*&)replaced);
        return replaced;
    }

    // array
    if (!shape.empty()) {
        int length = shape.back();
        shape.pop_back();
        auto symbol =
            Symbol::CreateArray(node->baseType, node->name, shape, length);

        if (node->init) {
            auto expr = node->init->As<InitializerExpr>();
            Assert(expr, "expect a initializer list");
            for (auto& [pos, init] : expr->parsed) {
                Visit(init);
                Cast(init, node->baseType);
            }
            if (symbolMap.GetContext() == Context::Global) {
                auto init = new std::unordered_map<int, ArithValue>();
                for (auto [pos, val] : expr->parsed) {
                    auto info = val->As<Expression>()->info->As<Constant>();
                    Assert(info, "expect constant in global variable");
                    (*init)[pos] = info->value;
                }
                symbol->info->As<Pointer>()->init = init;
            }
        }
        symbolMap.Add(symbol);
        block.back()->local.push_back(symbol);

        if (!node->init) return new VoidExpr();
        auto replaced = new AssignExpr();
        replaced->lhs = new IdentifierExpr(node->name);
        replaced->rhs = node->init;
        Visit((ASTNode*&)replaced);
        return replaced;
    }

    // value
    auto symbol = Symbol::CreateLValue(node->baseType, node->name);
    if (node->init) {
        Cast(node->init, node->baseType);
    }
    if (symbolMap.GetContext() == Context::Global) {
        if (!node->init)
            symbol->info->As<LValue>()->init = ArithValue(0);
        else {
            auto info = node->init->As<Expression>()->info->As<Constant>();
            Assert(info, "expect constant in global variable");
            symbol->info->As<LValue>()->init = info->value;
        }
    }
    symbolMap.Add(symbol);
    block.back()->local.push_back(symbol);

    if (symbolMap.GetContext() == Context::Local) {
        if (!node->init) return new VoidExpr();
        auto replaced = new AssignExpr();
        replaced->lhs = new IdentifierExpr(node->name);
        replaced->rhs = node->init;
        Visit((ASTNode*&)replaced);
        return replaced;
    }
    if (symbolMap.GetContext() == Context::Argument) {
        ASTNode* replaced = new IdentifierExpr(node->name);
        Visit(replaced);
        return replaced;
    }
    return node;
}

DECL_VISIT(FuncDecl) {
    retType = node->retType;
    auto symbol = Symbol::CreateFunction(retType, node->name, node);
    auto info = symbol->info->As<Function>();
    symbolMap.Add(symbol);

    symbolMap.CreateEnv(Context::Argument);
    for (auto& arg : node->args) {
        Visit(arg);
        auto argInfo = arg->As<Expression>()->info;
        Type type(argInfo->baseType, argInfo->As<Pointer>());
        info->argType.push_back(type);
    }

    Visit(node->content);
    symbolMap.ExitEnv();
    return node;
}

DECL_VISIT(IfStmt) {
    Visit(node->cond);
    if (node->tVal) Visit(node->tVal);
    if (node->fVal) Visit(node->fVal);
    return node;
}

DECL_VISIT(WhileStmt) {
    Visit(node->cond);
    Visit(node->content);
    return node;
}

DECL_VISIT(ContinueStmt) { return node; }

DECL_VISIT(BreakStmt) { return node; }

DECL_VISIT(ReturnStmt) {
    if (node->ret) Visit(node->ret);
    auto ret = node->ret->As<Expression>();
    Assert(ret, "return value should be an expression");
    if (!ret->info->As<Value>()) {
        node->ret = new LiteralExpr(0);
        Visit(node->ret);
    }
    if (retType != Type::Basic::Void) Cast(node->ret, retType);
    return node;
}

void SemanticAnalyzer::Visit(ASTNode*& node) {
    if (visited.count(node)) return;
    visited.insert(node);
    ASTVisitor::Visit(node);
}

void* SemanticAnalyzer::Run(void* input) {
    auto node = (AST::ASTNode*)input;
    Visit(node);
    return &symbolMap;
}
